# main.py

"""
OBJ模型体素化并转JSON主程序 (Open3D 单引擎、多材质版)
======================================================
(V7 - 移除硬编码版)
"""
import json
import os
import time
import argparse
from datetime import datetime, timezone
from typing import Dict, Any
from tqdm import tqdm
try:
    from . import obj_parser
    from . import voxelizer_open3d
    from .constants import EPSILON, DEFAULT_MATERIAL_NAME
except ImportError as e:
    print(f"错误：无法从包 'objTojson_ts' 导入必要的模块。错误信息: {e}")
    exit(1)

def float_color_to_rgba(kd_color: list[float] | None) -> list[int]:
    if kd_color and len(kd_color) == 3:
        return [int(c * 255) for c in kd_color] + [255]
    return [128, 128, 128, 255]

def process_obj_to_voxel_json(obj_file_path: str, output_json_path: str, voxel_size_setting: float, json_description: str, mode: str = 'solid') -> bool:
    obj_file_path = os.path.abspath(obj_file_path)
    output_json_path = os.path.abspath(output_json_path)
    print(f"\n--- 开始处理文件: {obj_file_path} ---")
    print(f"  指定体素大小: {voxel_size_setting}")
    print(f"  输出JSON到: {output_json_path}")
    is_solid_mode = (mode == 'solid')
    print(f"  选用体素化模式: {'实心 (Solid)' if is_solid_mode else '表面 (Surface)'}")
    process_start_time = time.time()
    print("步骤1: 加载OBJ模型数据...")
    model_data = obj_parser.load_obj(obj_file_path)
    if not model_data or not model_data.get('vertices'):
        return False
    material_library = model_data.get("materials", {})
    print("步骤2: 执行Open3D体素化...")
    voxelize_start_time = time.time()
    result = voxelizer_open3d.voxelize_with_open3d(model_data, voxel_size_setting, solid=is_solid_mode)
    voxelize_end_time = time.time()
    print(f"  引擎处理完成。耗时: {voxelize_end_time - voxelize_start_time:.4f} 秒。")
    if not result or result[0] is None:
        return False
    active_voxels_with_material, grid_origin, grid_dimensions = result
    voxelization_method_str = f"open3d_{'raycast_solid' if is_solid_mode else 'surface'}"
    print("步骤3: 构建动态调色板并组装JSON数据...")
    used_material_names = sorted(list(set(v['material'] for v in active_voxels_with_material)))
    voxel_palette = []
    material_to_id_map: Dict[str, int] = {}
    if not used_material_names or (len(used_material_names) == 1 and used_material_names[0] == DEFAULT_MATERIAL_NAME):
        used_material_names = [DEFAULT_MATERIAL_NAME]
    for i, mat_name in enumerate(used_material_names):
        palette_id = i + 1
        material_to_id_map[mat_name] = palette_id
        material_properties = material_library.get(mat_name, {})
        color_kd = material_properties.get("Kd")
        rgba_color = float_color_to_rgba(color_kd)
        voxel_palette.append({"id": palette_id, "name": mat_name, "color_rgba": rgba_color})
    if not voxel_palette:
        voxel_palette.append({"id": 1, "name": DEFAULT_MATERIAL_NAME, "color_rgba": [128, 128, 128, 255]})
        material_to_id_map[DEFAULT_MATERIAL_NAME] = 1
    lods_voxels_data = []
    progress_bar_desc = "  └ 组装JSON体素"
    for voxel in tqdm(active_voxels_with_material, desc=progress_bar_desc, ncols=80):
        mat_name = voxel.get('material', DEFAULT_MATERIAL_NAME)
        type_id = material_to_id_map.get(mat_name, 1)
        lods_voxels_data.append([int(voxel["x"]), int(voxel["y"]), int(voxel["z"]), type_id])
    json_data = {
      "format_version": "2.0-lod", "description": json_description,
      "metadata": { "source_file": os.path.basename(obj_file_path), "author": "My3DFrame OBJ Voxelizer (Multi-Material Edition)", "timestamp": datetime.now(timezone.utc).isoformat().replace('+00:00', 'Z'), "voxelization": { "method": voxelization_method_str, "source_voxel_size": voxel_size_setting, "parameters": {} } },
      "world_transform": { "grid_origin_meters": {"x": float(round(grid_origin[0], 6)), "y": float(round(grid_origin[1], 6)), "z": float(round(grid_origin[2], 6))}, "grid_dimensions_voxels": {"width": int(grid_dimensions[0]), "height": int(grid_dimensions[1]), "depth": int(grid_dimensions[2])} },
      "voxel_palette": voxel_palette,
      "lods": [{ "level": 0, "description": "最高细节级别。", "voxel_count": len(lods_voxels_data), "voxels": lods_voxels_data }],
      "voxel_attributes": []
    }
    try:
        output_dir = os.path.dirname(output_json_path)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        with open(output_json_path, 'w', encoding='utf-8') as json_file:
            json.dump(json_data, json_file, indent=2, ensure_ascii=False)
        print(f"  JSON文件写入成功！")
    except Exception as e:
        print(f"错误: 写入JSON文件失败: {e}")
        return False
    process_end_time = time.time()
    print(f"--- 文件 '{os.path.basename(obj_file_path)}' 处理完成。总耗时: {process_end_time - process_start_time:.4f} 秒 ---")
    return True

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="使用 Open3D 将OBJ模型文件体素化，并转换为 'My3DFrame v2.0-lod' 格式的JSON文件 (支持多材质)。", formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("input_obj", nargs='?', help="输入的OBJ文件路径。")
    parser.add_argument("-o", "--output_json", default="output_voxel.json", help="输出的JSON文件路径。")
    parser.add_argument("-s", "--voxel_size", type=float, default=1.0, help="体素大小 (世界单位)。")
    parser.add_argument("-d", "--description", default="Voxelized model", help="JSON文件中的描述信息")
    parser.add_argument('--version', action='version', version='%(prog)s 6.0-no-hardcoding')
    parser.add_argument("--mode", type=str, default="solid", choices=['solid', 'surface'], help="选择体素化模式:\n'solid':   生成实心模型 (默认)。\n'surface': 生成表面外壳模型。")
    args = parser.parse_args()
    if args.input_obj:
        if not os.path.exists(args.input_obj):
             print(f"错误: 输入的OBJ文件不存在: {args.input_obj}")
             exit(1)
        if args.voxel_size <= EPSILON:
            print(f"错误: 体素大小必须是正数。")
            exit(1)
        process_obj_to_voxel_json(args.input_obj, args.output_json, args.voxel_size, args.description, mode=args.mode)
    else:
        print("用法提示: 请提供一个输入OBJ文件。")
        parser.print_help()